#!/usr/bin/env python3
# encoding: utf-8
"""
@author: 孙伟伟
@contact: jakesun2020@163.com
@application: 以一个人名为输入, 使用模型帮助我们判断它最有可能是来自哪一个国家的人名, 这在某些国际化公司的业务中具有重要意义, 在用户注册过程中, 会根据用户填写的名字直接给他分配可能的国家或地区选项, 以及该国家或地区的国旗, 限制手机号码位数等等.
@file: person name.py
@time: 2021/11/20/020 21:38
"""
# 从io中导入文件打开方法
from io import open
# 帮助使用正则表达式进行子目录的查询
import glob
import os
# 用于获得常见字母及字符规范化
import string
import unicodedata

import random
import time
import math
import torch
# 导入nn准备构建模型
import torch.nn as nn
import matplotlib.pyplot as plt


'''
    第一步: 导入必备的工具包.
    第二步: 对data文件中的数据进行处理，满足训练要求.
    第三步: 构建RNN模型(包括传统RNN, LSTM以及GRU).
    第四步: 构建训练函数并进行训练.
    第五步: 构建评估函数并进行预测.
'''

'''1.对data进行处理'''
# 获取所有常用字符包括字母和常用标点
all_letters = string.ascii_letters + " .,;'"#
# 获取常用字符数量
n_letters = len(all_letters)# 57

# 字符规范化 将Unicode编码为ASCII
# 我们认为这个函数的作用就是去掉一些语言中的重音标记
# 如: Ślusàrski ---> Slusarski
def unicodeToAscii(s):
    return ''.join(
        c for c in unicodedata.normalize('NFD', s)
        if unicodedata.category(c) != 'Mn'
        and c in all_letters
    )

# 从硬盘中读取数据到内存
data_path = './data/names/'
def read_names(filename):
    # 从传入的文件中读取数据
    with open(filename, encoding='utf-8') as f:
        # 打开指定文件的内容，并使用strip取出开头和结尾的空格字符，然后用\n进行拆分
        names = f.read().strip().split('\n')# 如果不用split进行拆分，则遍历的结果为一个一个的字符文件
        return [unicodeToAscii(name) for name in names]
# filename = data_path + 'Chinese.txt'
# chinestname = read_names(filename)
# print(chinestname)

# 把各个国家的人名存储到一个字典中：{"English":["Lily", "Susan", "Kobe"], "Chinese":["Zhang San", "Xiao Ming"]}
# 把国家存储在一个列表中：["English",...,"Chinese"]
category_contry_names = {}# {"English":["Lily", "Susan", "Kobe"], "Chinese":["Zhang San", "Xiao Ming"]}
category_contry = []# ["English",...,"Chinese"]

# # 把所有人名数据进行遍历，使用glob.glob,因为可以进行全局正则
# for filename in glob.glob(data_path + '*.txt'):
#     # 每个文件名就是一个国家，文件内容为这个国家的人名
#     category = os.path.splitext(os.path.basename(filename))[0]
#     category_contry.append(category)
#     lines = read_names(filename)
#     category_contry_names[category] = lines


for filename in os.listdir(data_path):# 遍历data_path下的所有文件，filename表示每个国家的人名数据集的文件名，filename:Englist.txt/Chinese.txt
    if filename.endswith(".txt"):
        category_contry.append(filename[:-4])
    # contry_names = read_names(os.path.join(data_path, filename))  # 读取人名数据集，读取English.txt
    '''如果data_path = './data/names/'， 那么os.path.join(data_Path, '.txt') == data_path + '.txt' 返回True '''
    contry_names = read_names(data_path + filename)  # 读取人名数据集，读取English.txt
    category_contry_names[filename[:-4]] = contry_names

# 所有国家的名字中共有多少个单词进行排列组合
# all_name = []
# for value in category_contry_names.values():
#     all_name.append(value)
#
# n_letters = ''
# letter_ = ''
#
# for name in all_name[0]:# name表示每个名字，是一个字符串
#     for name_ in name:# 遍历每一个名字，生成单个字符
#         # + 把所有名字组成的字符进行拼接，即所有的名字组合成一个字符串。join是字符串合并：list_name = ['绮梦', '冷伊一', '香凝', '黛兰'] "*",join(list_name) 输出：绮梦*冷伊一*香凝*黛兰
#         n_letters = n_letters+ name_
# # print(n_letters)
# all_letters_ = set(n_letters)# 通过集合去除重复元素
#
# for letter in all_letters_:
#     letter_ = letter_ + letter
# all_letters = letter_
# n_letters = len(all_letters)# n_letters: 40
#
# print(all_letters)
# print(n_letters)

# 进入模型的参数都是tensor，因此需要把人名转化为tensor格式
def nameToOne_hot(name):# name为category_contry_names["Chinese"]对应的values(list)的其中一个数据
    # 每个人名有len(name)行，一行有n_letters个维度，
    tensor1 = torch.zeros(len(name), 1, n_letters)# 分为len(name)块，其中一块是维度是(1, n_letters),即每个人名用一个字母用1 * n_letters的张量表示
    for i, letter in enumerate(name):# 遍历每一个字母
        tensor1[i][0][all_letters.find(letter)] = 1
        '''str.find(str, beg=0)     str: 指定检索的字符串 beg: 开始索引，默认为0'''
    return  tensor1

# print(nameToOne_hot(category_contry_names["Chinese"][0]))
# 3.构建RNN模型
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        # input_size: RNN输入的最后一个维度
        # hidden_size: 隐藏层的最后一个维度
        # output_size: 最后线性层的输出维度
        # num_layer： RNN网络的层数
        super(RNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_layers = num_layers

        # 实例化预定义的RNN，参数分别为input_size, hidden_size, num_layers
        self.rnn = nn.RNN(input_size, hidden_size, num_layers)
        # 实例化全连接层，作用是将RNN的输出维度转换成指定维度的输出
        self.linear = nn.Linear(hidden_size, output_size)
        # 实例化nn中预定义的softmax，用于从输出层中获得类别结果
        self.softmax = nn.LogSoftmax(dim=-1)# dim=-1指定在最后一个维度上进行softmax

    def forward(self, input1, hidden):
        # input1: 表示人名分类器中的输入张量，形状是 1 * n_letters
        # hidden: 表示RNN的隐藏层张量，形状是 self.num_layer * 1 * self.hidden_size
        '''RNN 的输入是一个三维张量，又input1是一个二维张量，因此要输入RNN，需要给input1使用unsqueeze()扩展维度'''
        input = input1.unsqueeze(0)
        # 将参数传入到RNN中，如果num_layer=1,即单层网络，那么输出的结果rr和hn相等
        rr, hn = self.rnn(input, hidden)
        return self.softmax(self.linear(rr)), hn

    def initHidden(self):
        """初始化隐层张量,因为RNN输入的h0为0，因此实例化为一个全零的矩阵"""
        # 初始化一个（self.num_layers, 1, self.hidden_size）形状的0张量
        return torch.zeros(self.num_layers, 1, self.hidden_size)

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        # input_size: LSTM输入张量的最后一个维度
        # hidden_size: 隐藏层的最后一个维度
        # output_size: 最后线性层的输出维度
        # num_layer： 网络的层数
        super(LSTM, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_layers = num_layers

        self.lstm = nn.LSTM(input_size, hidden_size, num_layers)
        self.linear = nn.Linear(hidden_size, output_size)
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, input1, hidden, c):# 只有LSTM有一个细胞状态，RNN和GRU都是两个输入
        input = input1.unsqueeze(0)
        rr, (hn, c) = self.lstm(input, (hidden, c))
        return self.softmax(self.linear(rr)), hn, c

    def initHidden(self):
        """初始化函数不仅初始化hidden还要初始化细胞状态c, 它们形状相同"""
        c = hidden = torch.zeros(self.num_layers, 1, self.hidden_size)
        return hidden, c

class GRU(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        # input_size: GRU输入张量的最后一个维度
        # hidden_size: 隐藏层的最后一个维度
        # output_size: 最后线性层的输出维度
        # num_layer： 网络的层数
        super(GRU, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_layers = num_layers

        self.gru = nn.GRU(input_size, hidden_size, num_layers)
        self.linear = nn.Linear(hidden_size, output_size)
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, input1, hidden):
        input = input1.unsqueeze(0)
        rr, hn = self.gru(input, hidden)
        return self.softmax(self.linear(rr)), hn

    def initHidden(self):
        return torch.zeros(self.num_layers, 1, self.hidden_size)

# 3.1实例化参数进行测试
input_size = n_letters#输入张量的大小为输入张量的最后一个维度，又输入张量是一个onehot编码(len(name) * 1 * n_letters),名字的长度为块数，一块里面的大小为1 * 57的onehot编码
hidden_size = 128# 隐藏层的最后一个维度
output_size = len(category_contry)# 输出大小为线性层的输出维度，又因为线性层的输出为分类问题，因此output_size为人名国家的分类
num_layer = 1# 网络的层数

# input = category_contry_names["Chinese"][0]
# 假设用B作为RNN的首次输入，通过nameToOne_hot()转化为onehot张量，转化为的张量为3维
# eg:   name = [[],[],[]],name表示有三个字符组成，因为最外面的[]包含了三个[], 第一个[]表示名字的第一个字符,长度为n_letters,大小为1 * n_lettters
# 通过nameToOne_hot()转化的张量为三维，又RNN中执行了input1.unsqueeze()增加了一维，所以需要降低一个维度
input = nameToOne_hot("Z").squeeze(0)

# input1 = nameToOne_hot("B")
# input1 = nameToOne_hot(category_contry_names["Chinese"][0])

# 初始化一个三维的隐藏层0的张量，同样也是LSTM的初始细胞状态张量
hidden = c = torch.zeros(1, 1, hidden_size)
# print(input)
# print(input.size())
# print(input1)
# print(input1.size())

# 实例化模型
rnn = RNN(input_size, hidden_size, output_size)
lstm = LSTM(input_size, hidden_size, output_size)
gru = GRU(input_size, hidden_size, output_size)

# 调用模型
# rnn_output, hn = rnn(input, hidden)
# print("rnn_output:\n", rnn_output)
# lstm_output, hn, c = lstm(input, hidden, c)
# print("rnn_output:\n", lstm_output)
# gru_output, hn = gru(input, hidden)
# print("gru_output:\n", gru_output)

# 4.构建训练函数并进行训练
# 4.1从输出结果中获得名字所属的国家
def categoryFromOutput(output):
    top_category, top_index = output.topk(1)
    top_index = top_index[0].item()
    '''
        x = torch.tensor([1,2,3,4,5,6,7])
        top_category, top_index = x.topk(1)
        ——> top_index = tensor([6])
            top_index[0] = tensor(6)
            top_index[0].item() = 6
    '''
    return category_contry[top_index], top_index# 从返回的tensor格式的output中取出18个类别中最大的那一个类别的索引，然后对应到category_contry中

# 4.2随机生成训练数据
def randomTrainingExample():
    '''
    return：
        category: 随机选择的国家 eg: Chinese ,str类型
        category_contry_names_: 从随机选择的国家中随机选出一个名字，category_contry_names[category] ——>Chinese对应的名字中随机选择一个
        category_tensor: 把随机选择的国家对应列表的下标返回，并转换为tensor格式
        category_contry_names_tensor_: 把category_contry_names_转化为onehot向量
    '''
    # 从所有对应的国家中随机选择一个国家category
    category = random.choice(category_contry)
    # 从上一条程序随机选择的国家category中随机选择与之对应的名字
    category_contry_names_ = random.choice(category_contry_names[category])
    category_tensor = torch.tensor([category_contry.index(category)], dtype=torch.long)
    '''
        list1 = ['English', 'Chinese', 'Frech', 'German', 'Italian', 'Russion', 'Spanish']
        list1.index("Italian")——> 4
        列表.index(列表中的元素)，返回值：传入元素的下标
    '''
    category_contry_names_tensor_ = nameToOne_hot(category_contry_names_)
    return category, category_contry_names_, category_tensor, category_contry_names_tensor_

# for i in range(10):
#     category, name, category_tensor, line_tensor = randomTrainingExample()
#     print('category =', category, '/ name =', name, '/ category_tensor =', category_tensor)

# 4.3构建RNN训练函数
# 损失函数NLLLOSS，因为和RNN最后一层的nn.LogSotemax()逻辑匹配
criterion = nn.NLLLoss()

learning_rate = 0.005
def trainRNN(category_tensor, category_name_tensor):
    '''
    传入的值必须是tensor格式,分别为label和input
    :param category_tensor: 相当于训练的标签，即各个名字的所属的国家，例如Chinese，English
    :param category_name_tensor: 相当于各个训练数据，例如Bai、Zhao
    :return: RNN的最后一层输出output， 损失loss
    '''
    # 初始化RNN的第一个隐藏层张量
    hidden = rnn.initHidden()

    # 梯度归零
    rnn.zero_grad()

    # 把name中的一个一个字符传入到RNN模型中训练
    for i in range(category_name_tensor.size()[0]):
        output, hidden = rnn(category_name_tensor[i], hidden)

    # category_tensor是一个二维张量，因此要给output降维
    loss = criterion(output.squeeze(0), category_tensor)
    loss.backward()
    # 更新网络参数
    for p in rnn.parameters():
        p.data.add_(-learning_rate, p.grad.data)

    return  output, loss.item()

def trainLSTM(category_tensor, category_name_tensor):
    # 初始化隐藏层状态，初始化细胞状态
    hidden, c = lstm.initHidden()
    lstm.zero_grad()

    # 把name中的一个一个字符传入到RNN模型中训练
    for i in range(category_name_tensor.size()[0]):
        output, hidden, c = lstm(category_name_tensor[i], hidden, c)

    # category_tensor是一个二维张量，因此要给output降维
    loss = criterion(output.squeeze(0), category_tensor)
    loss.backward()
    # 更新网络参数
    for p in lstm.parameters():
        p.data.add_(-learning_rate, p.grad.data)

    return  output, loss.item()


def trainGRU(category_tensor, category_name_tensor):
    # 初始化隐藏层状态，初始化细胞状态
    hidden = gru.initHidden()
    gru.zero_grad()

    # 把name中的一个一个字符传入到RNN模型中训练
    for i in range(category_name_tensor.size()[0]):
        output, hidden = gru(category_name_tensor[i], hidden)

    # category_tensor是一个二维张量，因此要给output降维
    loss = criterion(output.squeeze(0), category_tensor)
    loss.backward()
    # 更新网络参数
    for p in gru.parameters():
        p.data.add_(-learning_rate, p.grad.data)

    return output, loss.item()

# 构建时间函数，用于查看训练了多久
def timeSince(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)# math.floor()向下取整        math.ceil()向上取整
    s -= m * 60
    return '%dm %ds' % (m ,s)

# 构建训练日志

# 设置训练迭代次数
n_iters = 100000
# 结果的打印间隔
print_every = 500
# 绘制损失函数曲线的间隔
plot_every = 1000

def train(train_type_fn):# train_type_fn表示模型的type：例如trainRNN

    # 每个损失值在图上绘制的值
    all_losses = []
    # 获取开始训练的时间戳
    start = time.time()
    # 设置初始间隔损失
    current_loss = 0
    # 从1开始训练，共进行n_liters次, 一个epoch有1000条数据，一个batch 20条数据， 那么训练一个epoch需要50次，那么这个50为iterations
    for iter in range(1, n_iters + 1):
        category, category_contry_names, category_tensor, category_contry_names_tensor = randomTrainingExample()
        '''
            category: 随机选择的国家 eg: Chinese ,str类型
            category_contry_names_: 从随机选择的国家中随机选出一个名字，category_contry_names[category] ——>Chinese对应的名字中随机选择一个
            category_tensor: 把随机选择的国家对应列表的下标返回，并转换为tensor格式
            category_contry_names_tensor_: 把category_contry_names_转化为onehot向量
        '''
        output, loss = train_type_fn(category_tensor, category_contry_names_tensor)
        current_loss += loss

        if iter % print_every == 0:# 如果迭代次数是50的倍数，那么打印结果
            predict, predict_index = categoryFromOutput(output)# 将output的输出张量中，取出最大可能性的那个国家类别和对应的下标
            correct = '✓' if predict == category else '✗ (%s)' % category
            # 打印迭代步, 迭代步百分比, 当前训练耗时, 损失, 该步预测的名字, 以及是否正确
            print('%d %d%% (%s) %.4f %s / %s %s' % (iter, iter / n_iters * 100, timeSince(start), loss, category, predict, correct))

        if iter % plot_every == 0 :# 把损失函数画出来
            all_losses.append(current_loss / plot_every)# 把要画的点存放到一个列表中，通过在列表中取值画图
            current_loss = 0

    return all_losses, int(time.time() - start)

# 开始训练
all_losses1, period1 = train(trainRNN)
all_losses2, period2 = train(trainLSTM)
all_losses3, period3 = train(trainGRU)

# 绘制损失对比曲线，训练耗时对比图
# 创建画图0，用于绘制损失值的变化，采用折线图
plt.figure(0)
plt.plot(all_losses1, label="RNN")
plt.plot(all_losses2, color="red", label="LSTM")
plt.plot(all_losses3, color="orange", label="GRU")

# 创建画图1，用于绘制训练耗时的对比，采用柱形图
plt.figure(1)
x_data = ["RNN", "LSTM", "GRU"]
y_data = [period1, period2, period3]
plt.bar(range(len(x_data)), y_data, tick_label=x_data)
plt.show()






















